Python performance for Spark